//
//  BinaryAVLManager.swift
//  BinaryTree
//
//  Created by 郭红旗 on 2021/11/21.
//

import UIKit

class BinaryAVLManager<E:Equatable>{
    typealias comparatorClosure = (_ e1:E,_ e2:E) -> Int
    var root:AVLTree<E>?
    private var comparator:comparatorClosure?
    /// 比较器
    init(comparator:@escaping comparatorClosure) {
        self.comparator = comparator
    }
    
    func add(element:E){
        if root == nil {
            root = createNode(element: element)
            afterAdd(node: root!)
            return
        }
        var node = root
        var lastNode = root
        var result:Int = 0
        while node != nil {
            result = comparator!(element,node!.value!)
            lastNode = node
            if result > 0 {
                node = node?.right
            }else if result < 0 {
                node  = node?.left
            }else{
                
            }
        }
        let newNode = createNode(element: element, parent: lastNode)
        if result > 0 {
            lastNode?.right = newNode
        }else if result < 0 {
            lastNode?.left = newNode
        }
        afterAdd(node: newNode)
    }
    
    func afterAdd(node:AVLTree<E>) {
        var tree:AVLTree<E>? = node
        while tree?.parent != nil {
            tree = tree?.parent
            if isBalance(node: tree!) {
                updateHight(node: tree!)
            }else{
                reBalance(node: tree!)
            }
        }
    }
    /// 恢复平衡
    func reBalance(node:AVLTree<E>) {
        let grandNode = node
        let parent = node.trailChild()
        let node = parent?.trailChild()
        if isLeft(node: parent) { // L
            if isLeft(node: node) { // ll
                rotateRight(grand: grandNode)
            }else{ // LR
                rotateLeft(grand: parent!)
                rotateRight(grand: grandNode)
            }
        }else if isRight(node: parent) { // R
            if isRight(node: node) { // RR
                rotateLeft(grand: grandNode)
            }else{ // RL
                rotateRight(grand: parent!)
                rotateLeft(grand: grandNode)
            }
        }
    }
    /// 右旋
    func rotateRight(grand:AVLTree<E>){
        guard let parent = grand.left else {
            return
        }
        grand.left = parent.right
        parent.right = grand
        parent.parent = grand.parent
        
        if isLeft(node: grand) {
            grand.parent?.left = parent
        }else if isRight(node: grand) {
            grand.parent?.right = parent
        }else{
            root = parent
        }
        grand.parent = parent
        
        updateHight(node: grand)
        updateHight(node: parent)
        
    }
    /// 左旋
    func rotateLeft(grand:AVLTree<E>){
        guard let parent = grand.right else {
            return
        }
        grand.right = parent.left
        parent.left = grand
        parent.parent = grand.parent
        if isLeft(node: grand) {
            grand.parent?.left = parent
        }
        if isRight(node: grand) {
            grand.parent?.right = parent
        }
        grand.parent = parent
        
        updateHight(node: grand)
        updateHight(node: parent)
    }
    /// 是否平衡
    func isBalance(node:AVLTree<E>) -> Bool {
        return node.balanceFactor()
    }
    /// 更新高度
    func updateHight(node:AVLTree<E>){
        node.updateHight()
    }
    /// 是否是父子树的左子树
    func isLeft(node:AVLTree<E>?) -> Bool {
        if node?.value == node?.parent?.left?.value {
            return true
        }
        return false
    }
    /// 是否是右子树
    func isRight(node:AVLTree<E>?) -> Bool {
        if node?.value == node?.parent?.right?.value {
            return true
        }
        return false
    }
    /// 创建节点
    func createNode(element:E,parent:AVLTree<E>? = nil) -> AVLTree<E> {
        return AVLTree.init(value: element, parent: parent)
    }
    
    

}


class AVLTree<T:Equatable>{
    var height:Int = 1
    var value:T?
    var left:AVLTree<T>?
    var right:AVLTree<T>?
    var parent:AVLTree<T>?
    /// 初始化
    init(value:T,parent:AVLTree<T>?) {
        self.value = value
        self.parent = parent
    }
    /// 打印函数
    var asString:String { return treeString(self){("\($0.value!)",$0.left,$0.right)}  }
    /// 平衡因子
    func balanceFactor() -> Bool {
        var leftHight = 0
        var rightHight = 0
        if left != nil {
            leftHight = left!.height
        }
        if right != nil {
            rightHight = right!.height
        }
        return abs(leftHight - rightHight) <= 1
    }
    /// 更新高度
    func updateHight(){
        var leftHight = 0
        var rightHight = 0
        if left != nil {
            leftHight = left!.height
        }
        if right != nil {
            rightHight = right!.height
        }
        height = 1 + max(leftHight, rightHight)
    }
    /// 获取左右子树中高度较高的子树
    func trailChild() -> AVLTree<T>? {
        var leftHight = 0
        var rightHight = 0
        if left != nil {
            leftHight = left!.height
        }
        if right != nil {
            rightHight = right!.height
        }
        if leftHight - rightHight > 0 {
            return left
        }
        if rightHight - leftHight > 0 {
            return right
        }
        return left
    }
    
}
